# rag/query.py
import numpy as np
from faiss_rag.embed import get_embedding
from faiss_rag.store import EntertainmentRAGStore
from loguru import logger
from typing import List,Dict,Any

def search_similar(query: str, top_k: int = 3) -> List[Dict]:
    store = EntertainmentRAGStore()
    if len(store) == 0:
        return []

    query_vec = np.array([get_embedding(query)], dtype=np.float32)
    distances, indices = store.index.search(query_vec, top_k)

    results = []
    for i, idx in enumerate(indices[0]):
        if idx < len(store.metadata):
            results.append({
                "text": store.metadata[idx]["text"],
                "score": float(distances[0][i])
            })
    return results

if __name__ == '__main__':
    results = search_similar('aaaaabbbbb',1000)
    # results = search_similar('aaaaabbbbb',5)
    logger.info(f"results是{results}")
